#!/bin/bash

eip_addr=$1
ban_action=$2
ban_port=$3

if [[ $eip_addr == "" ]] || [[ $ban_action == "" ]] || [[ ${ban_port//[0-9]/} != "" ]];
then
  echo "Usage: ./ban-eip-http.sh <eip_addr> <ban/unban> [port]"
  echo "       Port must be int, if not set, then default value is 80/443."
  exit 1
fi

python_env="export PYTHONPATH=/pitrix/lib/pitrix-common:/pitrix/lib/pitrix-cli:\${PYTHONPATH};"
cli_conf="/pitrix/conf/client.yaml"
exec_sql="/pitrix/bin/exec_sql.py"
webservice_host=$(grep webservice /etc/hosts|head -n 1|awk '{print $3}')

get_sg_id (){
  get_eip_id_sql="select eip_id from eip where eip_addr='${eip_addr}' and status='associated' and controller='self';"
  eip_id=$(${exec_sql} -d zone -c "${get_eip_id_sql}" 2>/dev/null|grep "eip-"|awk '{print $2}')
  if [[ $eip_id == "" ]];
  then
    echo "Warning: Associated eip ${eip_addr} is not found."
    exit 11
  fi

  resource_id=$(ssh -Tq $webservice_host "${python_env}/pitrix/cli/describe-eips -e ${eip_id} -f ${cli_conf}"|grep '"resource_id":"'|awk -F'"' '{print $4}')

  if [[ $(echo $resource_id|grep 'i-'|wc -l) == 1 ]];
  then
    get_sg_id_sql="select group_id from instance_security_group where instance_id='${resource_id}';"
    sg_id_zone=$(${exec_sql} -d zone -c "${get_sg_id_sql}" 2>/dev/null|grep "sg-"|awk '{print $2}')
  elif [[ $(echo $resource_id|grep -E 'lb-|rtr-'|wc -l) == 1 ]];
  then
    get_sg_id_sql="select group_id from resource_security_group where resource_id='${resource_id}';"
    sg_id_zone=$(${exec_sql} -d zone -c "${get_sg_id_sql}" 2>/dev/null|grep "sg-"|awk '{print $2}')
    if [[ $sg_id_zone == "" ]];
    then
      get_sg_id_sql="select group_id from instance_security_group where resource_id='${resource_id}';"
      sg_id_zone=$(${exec_sql} -d zone -c "${get_sg_id_sql}" 2>/dev/null|grep "sg-"|awk '{print $2}')
    fi
  else
    echo "Warning: ${eip_id}'s resource_id is not found."
    exit 12
  fi

  if [[ $sg_id_zone == "" ]];
  then
    echo "Warning: security group is None, [resource_id: ${resource_id}]."
    exit 12
  fi
}


add_ban_rule (){
  if [[ $2 == "" ]];
  then
    ban80_rule_id=$(ssh -Tq $webservice_host "${python_env}/pitrix/cli/add-security-group-rules -s $1 -r '[{\"protocol\":\"tcp\",\"priority\":\"0\",\"action\":\"drop\",\"val1\":\"80\",\"val2\":\"80\"}]' -f ${cli_conf}"|grep 'sgr-'|awk -F'"' '{print $2}')
    ban443_rule_id=$(ssh -Tq $webservice_host "${python_env}/pitrix/cli/add-security-group-rules -s $1 -r '[{\"protocol\":\"tcp\",\"priority\":\"0\",\"action\":\"drop\",\"val1\":\"443\",\"val2\":\"443\"}]' -f ${cli_conf}"|grep 'sgr-'|awk -F'"' '{print $2}')
    if [[ $ban80_rule_id == "" ]] || [[ $ban443_rule_id == "" ]];
    then
      echo "Warning: Failed to add 80/443 security group rules, [security_group_id: $1]."
      exit 31
    fi
    update80_sql="update security_group_rules set controller='pitrix' where group_rule_id='${ban80_rule_id}';"
    update443_sql="update security_group_rules set controller='pitrix' where group_rule_id='${ban443_rule_id}';"
    affected80_count=$(${exec_sql} -d zone -c "${update80_sql}" 2>/dev/null|grep 'affected count: 1'|wc -l)
    affected443_count=$(${exec_sql} -d zone -c "${update443_sql}" 2>/dev/null|grep 'affected count: 1'|wc -l)
    if [[ $affected80_count != 1 ]] || [[ $affected443_count != 1 ]];
    then
      echo "Error: Failed to update ${ban80_rule_id}/${ban443_rule_id} controller."
      exit 32
    fi
  else
    ban_rule_id=$(ssh -Tq $webservice_host "${python_env}/pitrix/cli/add-security-group-rules -s $1 -r '[{\"protocol\":\"tcp\",\"priority\":\"0\",\"action\":\"drop\",\"val1\":\"$2\",\"val2\":\"$2\"}]' -f ${cli_conf}"|grep 'sgr-'|awk -F'"' '{print $2}')
    if [[ $ban_rule_id == "" ]];
    then
      echo "Warning: Failed to add $2 security group rules, [security_group_id: $1]."
      exit 31
    fi
    update_sql="update security_group_rules set controller='pitrix' where group_rule_id='${ban_rule_id}';"
    affected_count=$(${exec_sql} -d zone -c "${update_sql}" 2>/dev/null|grep 'affected count: 1'|wc -l)
    if [[ $affected_count != 1 ]];
    then
      echo "Error: Failed to update ${ban_rule_id} controller."
      exit 32
    fi
  fi

  apply_result=$(ssh -Tq $webservice_host "${python_env}/pitrix/cli/apply-security-group -s $1 -f ${cli_conf}"|grep '"ret_code":0'|wc -l)
  if [[ $apply_result != 1 ]];
  then
    echo "Error: Failed to apply security group ${1}."
    exit 33
  fi
}


del_ban_rule (){
  if [[ $2 == "" ]];
  then
    get_sgr_id_sql="select group_rule_id from security_group_rules where group_id='$1' and controller='pitrix' and action='drop' and (val1 in ('80','443') or val2 in ('80','443'));"
    sgr_id_zone=$(${exec_sql} -d zone -c "${get_sgr_id_sql}" 2>/dev/null|grep 'sgr-'|awk '{print $2}')
  else
    get_sgr_id_sql="select group_rule_id from security_group_rules where group_id='$1' and controller='pitrix' and action='drop' and (val1='$2' or val2='$2');"
    sgr_id_zone=$(${exec_sql} -d zone -c "${get_sgr_id_sql}" 2>/dev/null|grep 'sgr-'|awk '{print $2}')
  fi

  if [[ $sgr_id_zone == "" ]];
  then
    echo "Warning: Get ban rule failed, maybe $eip_addr has no ban rule, [security_group_id: $1]."
    exit 41
  fi

  for sgr_id in $sgr_id_zone;
  do
    del_rule_result=$(ssh -Tq $webservice_host "${python_env}/pitrix/cli/delete-security-group-rules -r $sgr_id -f ${cli_conf}"|grep '"ret_code":0'|wc -l)
    if [[ $del_rule_result != 1 ]];
    then
      echo "Error: Failed to delete security group ${1} rule, [security_group_rule_id: $sgr_id]."
      exit 42
    fi
  done

  del_apply_result=$(ssh -Tq $webservice_host "${python_env}/pitrix/cli/apply-security-group -s $1 -f ${cli_conf}"|grep '"ret_code":0'|wc -l)
  if [[ $del_apply_result != 1 ]];
  then
    echo "Error: Failed to apply security group ${1}."
    exit 43
  fi
}


check_rule (){
  if [[ $2 == "" ]];
  then
    check_sgr_id_sql="select group_rule_id from security_group_rules where group_id='$1' and protocol='tcp' and priority=0 and direction=0 and (val1='443' or val1='80' or val2='443' or val2='80');"
    check_sgr_id_zone=$(${exec_sql} -d zone -c "${check_sgr_id_sql}" 2>/dev/null|grep 'sgr-'|awk '{print $2}')
  else
    check_sgr_id_sql="select group_rule_id from security_group_rules where group_id='$1' and protocol='tcp' and priority=0 and direction=0 and (val1='$2'or val2='$2');"
    check_sgr_id_zone=$(${exec_sql} -d zone -c "${check_sgr_id_sql}" 2>/dev/null|grep 'sgr-'|awk '{print $2}')
  fi

  if [[ $check_sgr_id_zone != "" ]];
  then
    for check_sgr_id in $check_sgr_id_zone;
    do
      modify_rule_result=$(ssh -Tq $webservice_host "${python_env}/pitrix/cli/modify-security-group-rule-attributes -r $check_sgr_id -p 1 -f ${cli_conf}"|grep '"ret_code":0'|wc -l)
      if [[ $modify_rule_result != 1 ]];
      then
        echo "Warning: Failed to modify security group $1 rule, [security_group_rule_id: ${check_sgr_id}]."
        exit 52
      fi
    done
  fi
}


get_sg_id

ban_date=$(date +%Y%m%d-%H%M%S)
shell_pwd=$(cd "$( dirname "$0" )" && pwd)
if [ ! -d "${shell_pwd}/logs" ];
then
  mkdir ${shell_pwd}/logs
fi
log_path="${shell_pwd}/logs/ban_eip_http-$(date +%Y%m).log"

if [[ $ban_action == "ban" ]];
then
  for sg_id in ${sg_id_zone};
  do
    check_rule $sg_id $ban_port
    add_ban_rule $sg_id $ban_port
  done
  if [[ $ban_port == "" ]];
  then
    echo "$ban_date #-# EIP $eip_addr ban 80/443 successful, [security_group_id: ${sg_id_zone}]."
    echo "$ban_date #-# EIP $eip_addr ban 80/443 successful, [security_group_id: ${sg_id_zone}]." >>${log_path}
  else
    echo "$ban_date #-# EIP $eip_addr ban $ban_port successful, [security_group_id: ${sg_id_zone}]."
    echo "$ban_date #-# EIP $eip_addr ban $ban_port successful, [security_group_id: ${sg_id_zone}]." >>${log_path}
  fi
  exit 0
elif [[ $ban_action == "unban" ]];
then
  for sg_id in ${sg_id_zone};
  do
    del_ban_rule $sg_id $ban_port
  done
  if [[ $ban_port == "" ]];
  then
    echo "$ban_date #-# EIP $eip_addr unban 80/443 successful, [security_group_id: ${sg_id_zone}]."
    echo "$ban_date #-# EIP $eip_addr unban 80/443 successful, [security_group_id: ${sg_id_zone}]." >>${log_path}
  else
    echo "$ban_date #-# EIP $eip_addr unban $ban_port successful, [security_group_id: ${sg_id_zone}]."
    echo "$ban_date #-# EIP $eip_addr unban $ban_port successful, [security_group_id: ${sg_id_zone}]." >>${log_path}
  fi
  exit 0
else
  echo "Usage: ./ban-eip-http.sh <eip_addr> <ban/unban> [port]"
  echo "       Port must be int, if not set, then default value is 80/443."
  exit 1
fi
